;  Vectorized Runge-Kutta
function agage_box_model_solver_RK4, chi, F, dt=dt

  A=F##chi
  B=F##(chi + dt*A/2.d)
  C=F##(chi + dt*B/2.d)
  D=F##(chi + dt*C)
  chi+=dt/6.d*(A + 2.d*B + 2.d*C + D)

  return, chi

end


;  Lorenz 4-cycle solver. 
;  Lorenz, E. N. (1971). AN N-CYCLE TIME-DIFFERENCING SCHEME FOR 
;  STEPWISE NUMERICAL INTEGRATION. Monthly Weather Review, 99(8), 644-648
function agage_box_model_solver, chi, F, dt=dt

  N=4.d
  delT=dt/N
  
  y=chi
  
  z=dblarr(12)
  
  a=0.d
  b=1.d/delT

  for i=0, N-1 do begin
    z=z*a

    z+=F##y

    z=z/b
    y+=z
    
    if a eq 0. then begin
      a-=1.d/delT
      b-=1.d/delT     
    endif
    
    a+=1.d/dT
    b+=1.d/dT
    
  endfor

  return, y

end

; docformat = 'rst'
;
;  AGAGE_BOX_MODEL
;
;+
;
; :Purpose:
;   Calculate semi-hemispheric monthly average mole fractions, based on given emissions, 
;   stratospheric and oceanic lifetimes and OH reaction rates.
;   
; :Inputs::
; 
;   ic: (Required) a 12-element array specifying initial conditions in each of the 12 boxes.
;     Units are pmol/mol (~parts per trillion by volume).
;     Boxes are ordered from north to south and then vertically::
;       0:3 (1000hPa - 500hPa) -> 90N-3ON, 30N-0N, 0N-30S, 30S-90S
;       4:7 (500hPa - 200hPa) -> 90N-3ON, 30N-0N, 0N-30S, 30S-90S
;       8:11 (200hPa - 0hPa) -> 90N-3ON, 30N-0N, 0N-30S, 30S-90S
;
;   q: (Required) a 4 x nMonths array of emissions into the four surface boxes.
;     Units are Gg/yr
;     The length of the simulation is determined by the number of months specified in the second dimension.
;   
; :Keywords:
; 
;   mol_mass: (Required) Molecular mass
;     Units=g/mol
;   arr_OH: (Optional) two-element array of Arrhenius bi-molecular, temperature-dependent reaction rates with OH.
;     First element=A, second element=-E/R
;   lifeTime: (Optional) 12 x nSeasons array of non-OH first-order lifetimes for each season in each box. 
;     E.g. to specify stratospheric lifetimes, enter lifetimes in elements 8-11 (first index).
;     If this keyword is not set (default) then a very large lifetime is set in each box (not including OH reactions).
;     Units=years
;   
;   OH_scale (Optional): 8 x nMonths array of scaling to apply to the OH field 
;   V1_scale (Optional): 8 x nMonths array of scaling to apply to the velocity timescales
;   T_scale (Optional): 17 x nMonths array of scaling to apply to mixing timescales
;   transport_input_extension (Optional): String filename extension if alternative transport parameters are required
;   startT (Optional): time index at which to restart simulation (requires in_restart to be supplied)
;   endT (Optional): time index at which to end restarted simulation
;   in_restart: instantaneous mass in each box at the start of each month. Simulation restarts from ith index, 
;     given by startT
;   calendar_year (Required if inter-annually varying OH field and/or transport parameter file specified): Initial calendar
;     year for simulation. Only required if inter-annually varying OH field and/or transport parameter file.
;   
;   ascii_output: (Optional) output to ascii file: ascii_output/mf.csv
;   all: (Optional) output all 12 boxes, rather than just 4 surface boxes.  
;     Output is then a 12 x nMonths array of mole fractions
;   overall_lifeTime (Optional, output): if variable is supplied, it will contain the monthly-mean overall lifetime.
;     Units=years
;   output_lifetime_boxes (optional, output): if specified, will determine lifetime based on loss rate from
;     a sub-set of all boxes. E.g. if output_lifetime_boxes=[8, 9, 10, 11] then stratospheric lifetime is output 
;     in overall_lifetime
;   average_lifeTime (Optional, output): average overall lifetime over the whole simulation.
;     Units=years
;   burden (Optional, output): if variable is supplied, it will contain the monthly-average global burden.
;     Units=g
;   loss (Optional, output): monthly average mass loss.
;     Units=g
;   emissions (Optional, output): monthly average mass emissions (used for diagnostics).
;     Units=g
;   restart (optional, output): will contain the instantaneous mass in each box at the start of each month
;     for use in subsequent simulations using the in_restart and startT keywords.
;     Units=g
;   
; :Outputs:
; 
;   4 x nMonths array of surface mole fractions is returned (pmol/mol).
;   If ALL keyword is set, a 12 x nMonths array of mole fractions from all boxes is returned (pmol/mol).
;
; :Example:
; 
;   Two year simulation::
;     q = replicate(1., 4, 24) ;Two year simulation, 1Gg/yr in all surface boxes
;     q[0, *]=10.  ;10Gg/yr emissions in northern box
;     ic = replicate(10., 12) ;Initially uniform mole fraction (10 pmol/mol).
;     ppt_out = agage_box_model(ic, q, mol_mass=50.)
;     print, ppt_out
;     
;       10.313749       10.060051       10.015412       10.047651
;       10.722616       10.217531       10.067380       10.109805
;       11.024041       10.418397       10.142118       10.164657
;       11.284635       10.609171       10.226286       10.206647...
;
;
; :Requires:
; 
;   IDL .sav files containing model parameters::
;     agage_box_model_input_parameters.sav: transport parameters
;     agage_box_model_input_OH.sav: OH concentrations
;     agage_box_model_input_temp.sav: temperatures
;
; :History:
;   Written by: Matt Rigby, University of Bristol, Feb 20, 2012
;
;-
function agage_box_model, ic, q, $
  ;Required keywords
  mol_mass=mol_mass, $
  ;Optional input keywords
  arr_OH=arr_OH, lifeTime=lifeTime, OH_scale=OH_scale, $
  in_restart=in_restart, startT=startT, endT=endT, IAV=IAV, $
  V1_scale=V1_scale, T_scale=T_scale, $
  transport_input_Extension=transport_input_Extension, OH_input_extension=OH_input_extension, $
  calendar_year=calendar_year, $
  OH_field=OH_field, temperature_field=temperature_field, transport_field=transport_field, $
  ;Output keywords
  overall_lifeTime=overall_lifeTime, output_lifetime_boxes=output_lifetime_boxes, $
  output_oh_lifetime=output_oh_lifetime, $
  average_lifetime=average_lifetime, $
  burden=burden, loss=loss, emissions=emissions, $
  all=all, restart=restart, ascii_output=ascii_output

  compile_opt idl2, hidden
;  on_error, 2

  @agage_box_model_path

  ;Parameters
  dt=2.d*24.d*3600.d	;seconds/timestep
  yearsToSec=365.25d*24.d*3600.d	;seconds/year
  nBox=12

  mAtm=5.1170001e+18  ;Mass of the atmosphere in kg
  mAir=28.97d ;dry molecular mass of air in g/mol
  mass=[replicate(0.25d*0.5d, 4), replicate(0.25d*0.3d, 4), replicate(0.2d*0.25d, 4)]*Matm*1000.d ;Mass of air in each box (g)


  ;Check inputs
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  if keyword_set(IAV) eq 0 then IAV=0
  
  if n_elements(q[*, 0]) ne 4 then begin
    message, 'Must be 4 surface boxes'
  endif
  nMonths=n_elements(q[0, *])
  nYears=nMonths/12  ;MONTHLY EMISSIONS ONLY, units: Gg/yr


  ;OH reaction rates (if not set, assume almost inert)
  if keyword_set(arr_OH) eq 0 then begin
    arrA=1.e-30
    arrER=1800.d
  endif else begin
    arrA=arr_OH[0]
    arrER=arr_OH[1]
  endelse

  
  ;If lifetimes are not specified, assume very large
  if n_elements(lifeTime) eq 0 then begin
    lifeTime=double(replicate(1.e12, 12, 4*nYears))
  endif
  
  ;Assume le 0 lifetime is a mistake and replace with very long lifetime
  if total(lifetime le 0.) gt 0 then begin
    lifetime[where(lifetime le 0.)]=1.e12
  endif



  ;Get OH field (default: Spivakovsky et al, 2000)
  if keyword_set(OH_field) eq 0 then begin
    if keyword_set(OH_input_extension) then begin
      OH_Fname=file_search(strcompress(agage_box_model_path + $
        'agage_box_model_input_oh_' + OH_input_extension + '.sav', /remove_all))  ;Contains OH[box, month]
      if OH_fname eq '' then begin
        message, "Can't find OH file"
      endif else begin
        restore, OH_fname
      endelse
      if n_elements(OH[0, *]) gt 12 then begin
        ;Inter-annually varying OH
        if n_elements(calendar_year) eq 0 then begin
          message, "Looks like you're trying to specify an input OH file. Make sure you also specify a valid start year."
        endif
        wh=where(OH_time ge float(calendar_year), count)
        if count lt nMonths then begin
          message, "Check that input OH file covers entire date range"
        endif
        OH=OH[*, wh[0]:wh[0] + nMonths-1]
      endif else begin
        ;Inter-annually repeating OH
        OH_temp=dblarr(12, nMonths)
        for yi=0, nYears-1 do begin
          OH_temp[*, yi*12:(yi+1)*12-1]=OH
        endfor
        OH=temporary(OH_temp)
      endelse
    endif else begin
      ;Get default OH file
      restore, strcompress(agage_box_model_path + 'agage_box_model_input_oh.sav', /remove_all)
      OH_temp=dblarr(12, nMonths)
      for yi=0, nYears-1 do begin
        OH_temp[*, yi*12:(yi+1)*12-1]=OH
      endfor
      OH=temporary(OH_temp)
    endelse

    ;If OH scaling not specified, use default
    OH_adj=replicate(1.d, 12, nMonths)
    if keyword_set(OH_scale) then begin
      oh_Adj[0:7, *]=OH_scale
    endif
    OH=OH*OH_adj
  endif else begin
    OH=OH_field
  endelse

  ;Get transport parameters (default: Spivakovsky et al, 2000)
  if keyword_set(transport_field) eq 0 then begin
    if keyword_set(transport_input_extension) then begin
      transport_Fname=file_search(strcompress(agage_box_model_path + $
        'agage_box_model_input_parameters_' + transport_input_extension + '.sav', /remove_all))
      if transport_fname eq '' then begin
        message, "Can't find transport file"
      endif else begin
        restore, transport_fname
      endelse
      
      if n_elements(T[0, *]) gt 12 then begin
        ;Restored parameters have inter-annual variability
        if n_elements(calendar_year) eq 0 then begin
          message, "Looks like you're trying to specify an input transport parameter file. Make sure you also specify a valid start year."
        endif
        wh=where(transport_parameter_time ge float(calendar_year), count)
        if count lt nMonths then begin
          message, "Check that input transport file covers entire date range"
        endif
        T=T[*, wh[0]:wh[0] + nMonths-1]
      endif else begin
        T_temp=dblarr(17, nMonths)
        for yi=0, nYears-1 do begin
          T_temp[*, yi*12:(yi+1)*12-1]=T
        endfor
        T=temporary(T_temp)        
      endelse
    endif else begin
      restore, strcompress(agage_box_model_path + 'agage_box_model_input_parameters.sav', /remove_all)
      T_temp=dblarr(17, nMonths)
      for yi=0, nYears-1 do begin
        T_temp[*, yi*12:(yi+1)*12-1]=T
      endfor
      T=temporary(T_temp)
    endelse
  endif else begin
    restore, strcompress(agage_box_model_path + 'agage_box_model_input_parameters.sav', /remove_all)
    T=transport_field
  endelse


  V1=V1*24.d*3600.d ;convert from days to seconds
  T=T*24.d*3600.d ;convert from days to seconds


  ;Get Temperature field (NCEP 1990 - 2008 monthly mean)
  if keyword_set(temperature_field) eq 0 then begin
    restore, strcompress(agage_box_model_path + 'agage_box_model_input_temp.sav', /remove_all)  ;Contains: Temp[box, month]
  endif else begin
    temp=temperature_field
  endelse


  ;If transport parameter scaling not specified, use default
  if keyword_set(T_scale) eq 0 then begin
    T_scale=replicate(1., n_elements(T[*, 0]), 12*nYears)
  endif

  ;Calculate and store transport matrices
  F=fltarr(12, 12, nMonths)
  if IAV then begin
    for mi=0, nMonths-1 do begin
      F[*, *, mi]=agage_box_model_transport(V1[*, (mi mod 12)], intersection_V1, $
      	T[*, mi]/T_scale[*, mi], intersection_T)
    endfor
  endif else begin
    for mi=0, 11 do begin
      F[*, *, mi]=agage_box_model_transport(V1[*, mi], intersection_V1, $
        T[*, mi]/T_scale[*, mi], intersection_T)
    endfor
  	for yi=1, nYears-1 do begin
  	  F[*, *, yi*12:(yi+1)*12-1]=F[*, *, 0:11]
  	endfor
  endelse

  ;Emissions
  qModel=dblarr(nBox, nYears*12L)
  qModel[0:3, *]=Q*1.e9/365.25d/24.d/3600.d  ;g/s
  
  TimeSize=long(float(nYears)*360.d*24.d*3600.d/dt)

  ;Declare output arrays
  cMonth=dblarr(nBox, nYears*12L)
  countMonth=dblarr(nBox, nYears*12L)
  overall_lifeTime=dblarr(nYears*12L)
  oh_lifeTime=dblarr(nYears*12L)
  burden=dblarr(nBox, nYears*12L)
  loss=dblarr(nYears*12L)
  emissions=dblarr(nYears*12L)
  average_lifetime=0.d

  yearRec=fltarr(timesize)
  monthRec=fltarr(timesize)

  restart=fltarr(12, nYears*12L)

  if n_elements(output_lifetime_boxes) eq 0 then begin
    lifetime_boxes=indgen(12)
  endif else begin
    lifetime_boxes=output_lifetime_boxes
  endelse

  month=-1

  ;Run model
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  ;Set initial conditions
  ;If restart supplied (with startT and endT), use that, otherwise use IC
  if n_elements(startT) gt 0 then begin
    if startT gt 0 then begin
      c=reform(in_restart[*, startT])
    endif else begin
      c=IC*mol_mass[0]/mAir*1.e-12*mass ;ppt to g
    endelse
    startMT=startT*30.d*24.d*3600.d/dt
    endMT=(endT+1)*30.d*24.d*3600.d/dt-1
  endif else begin
    startMT=0
    endMT=TimeSize-1
    c=IC*mol_mass[0]/mAir*1.e-12*mass ;ppt to g
  endelse

  ;Start time stepping
  for ti=startMT, endMT do begin

    ;Time
    year=long(ti*dt/(360.*24.*3600.))
    monthPrev=month
    month=long( ( ti*dt - float(year*360.*24.*3600.) )/30./24./3600.)
    season=month/3

    ;Save restart at start of month change
    if month ne monthPrev then restart[*, month + 12*year]=c

    ;Stratospheric lifetime and oceanic lifetimes for current SEASON
    lifetime_ti=lifeTime[*, year*4 + season]

    ;Step forward RK4 solver by default (Transport only. Loss and emissions left out of this step for efficiency)
    ;Delete "_RK4" for Lorentz solver
    c=agage_box_model_solver_RK4(c/mass, reform(F[*, *, 12*year + month]), dt=dt)*mass

    ;Loss and emissions
    oh_loss=c*OH[*, 12*year + month]*arrA*exp(arrER/Temp[*, month])*dt*365.d/360.d  
    loss_ti=c/lifetime_ti/yearsToSec*dt*365.d/360.d ;oh_loss + 

    emissions_ti=qModel[*, year*12 + month]*dt*365.d/360.d
    c += emissions_ti - loss_ti - oh_loss

    ;Monthly averages
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    
    cMonth[*, 12*year + month]+=c
    
    if total(oh_loss) eq 0. then begin
      oh_loss=replicate(1.e-24, 12)
    endif
    if total(loss_ti[lifetime_boxes]) eq 0. then begin
      loss_ti[lifetime_boxes]=replicate(1.e-24, n_elements(lifetime_boxes))
    endif

    oh_lifetime[12*year + month]+=total(c)/total(oh_loss)*dt
    overall_lifeTime[12*year + month]+=total(c)/total(loss_ti[lifetime_boxes])*dt

    burden[*, 12*year + month]+=c

    if total(loss_ti) eq 0. then begin
      loss_ti=replicate(1.e-24, 12)
    endif
    loss[12*year + month]+=total(loss_ti)/dt

    emissions[12*year + month]+=total(emissions_ti)/dt
    countMonth[*, 12*year + month]++
    
    average_lifetime+=total(c)/total(loss_ti[lifetime_boxes])*dt

  endfor

  ;Convert to ppt and calculate averages
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  
  wh=where(countMonth gt 0)
  cMonth[wh]=cMonth[wh]/double(countMonth[wh])/mol_mass[0]*mAir*1.e12
  burden[wh]=burden[wh]/double(countMonth[wh])
  for bi=0, nBox-1 do begin
    cMonth[bi, *]=cMonth[bi, *]/mass[bi]
  endfor

  wh=where(countMonth[0, *] gt 0, count)
  overall_lifeTime[wh]=overall_lifeTime[wh]/double(countMonth[0, wh])*365.d/360.d
  if keyword_set(output_oh_lifetime) then begin
    overall_lifeTime[wh]=oh_lifeTime[wh]/double(countMonth[0, wh])*365.d/360.d
  endif
  loss[wh]=loss[wh]/double(countMonth[0, wh])
  emissions[wh]=emissions[wh]/double(countMonth[0, wh])*360.d/365.d

  average_lifetime=average_lifetime/double(endMT - startMT)*365.d/360.d

  ;Prepare output
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  
  if keyword_set(all) eq 0 then begin
    cMonth=cMonth[0:3, *]
    burden=total(burden, 1)
  endif

  ;Write output file if required
  if keyword_set(ascii_output) then begin
    if !version.os_family eq 'Windows' then slash='\' else slash='/'
    openw, fun, strcompress(agage_box_model_path + 'ascii_output' + slash + 'mf.csv', /remove_all), /get_lun
      if keyword_set(all) then begin
        printf, fun, 'Box 0, Box 1, Box 2, Box 3, Box 4, Box 5, Box 6, Box 7, Box 8, Box 9, Box 10, Box 11,'
        for month=0, n_elements(cMonth[0, *])-1 do begin
          printf, fun, cMonth[*, month], format="(12(f16.4, ','))"
        endfor
      endif else begin
        printf, fun, 'Box 0, Box 1, Box 2, Box 3,'
        for month=0, n_elements(cMonth[0, *])-1 do begin
          printf, fun, cMonth[*, month], format="(4(f16.4, ','))"
        endfor        
      endelse
    close, fun
    free_lun, fun
  endif

  ;Return monthly means
  return, cMonth

end
